"""
Boundary Conditions Module

Handles boundary condition application for multi-domain Navier-Stokes
simulations with spectral element methods. Supports Dirichlet, Neumann,
and interface coupling conditions.
"""

"""
    apply_multidomain_boundary_conditions!(u::Array, v::Array, w::Union{Array,Nothing}, 
                                           p::Array, multidomain::MultiDomain{D}, 
                                           x::Vector, y::Vector, z::Union{Vector,Nothing}=nothing;
                                           U_lid::Float64=1.0) where D

Apply boundary conditions for multi-domain Navier-Stokes simulation.
"""
function apply_multidomain_boundary_conditions!(u::Array, v::Array, w::Union{Array,Nothing}, 
                                                p::Array, multidomain::MultiDomain{D}, 
                                                x::Vector, y::Vector, z::Union{Vector,Nothing}=nothing;
                                                U_lid::Float64=1.0) where D
    
    # Apply velocity boundary conditions
    apply_velocity_boundary_conditions!(u, v, w, multidomain, x, y, z, U_lid)
    
    # Apply pressure boundary conditions  
    apply_pressure_boundary_conditions!(p, multidomain, x, y, z)
    
    # Apply interface coupling conditions
    apply_interface_coupling!(u, v, w, p, multidomain)
    
    return nothing
end

"""
    apply_velocity_boundary_conditions!(u::Array, v::Array, w::Union{Array,Nothing}, 
                                       multidomain::MultiDomain{D}, x::Vector, y::Vector, 
                                       z::Union{Vector,Nothing}, U_lid::Float64) where D

Apply velocity boundary conditions for each domain block.
"""
function apply_velocity_boundary_conditions!(u::Array, v::Array, w::Union{Array,Nothing}, 
                                            multidomain::MultiDomain{D}, x::Vector, y::Vector, 
                                            z::Union{Vector,Nothing}, U_lid::Float64) where D
    
    # Process each active block
    for l in 1:multidomain.index_sum
        if D == 2
            apply_2d_velocity_bc!(u[l, :, :], v[l, :, :], multidomain, l, x, y, U_lid)
        else
            apply_3d_velocity_bc!(u[l, :, :, :], v[l, :, :, :], w[l, :, :, :], 
                                 multidomain, l, x, y, z, U_lid)
        end
    end
end

"""
    apply_2d_velocity_bc!(u_block::AbstractMatrix, v_block::AbstractMatrix, 
                         multidomain::MultiDomain{2}, block_idx::Int, 
                         x::Vector, y::Vector, U_lid::Float64)

Apply 2D velocity boundary conditions for a single block.
"""
function apply_2d_velocity_bc!(u_block::AbstractMatrix, v_block::AbstractMatrix, 
                              multidomain::MultiDomain{2}, block_idx::Int, 
                              x::Vector, y::Vector, U_lid::Float64)
    
    # Get block coordinates in the distribution matrix
    block_coord = multidomain.g[block_idx]
    i0, i1 = block_coord[1], block_coord[2]
    
    nx, ny = size(u_block)
    
    # Apply boundary conditions based on block position and neighboring blocks
    
    # Left boundary (x = x[1])
    if multidomain.matrix_dist[i0-1, i1] == 0  # External boundary
        u_block[1, :] .= 0.0  # No-slip wall
        v_block[1, :] .= 0.0
    end
    
    # Right boundary (x = x[end])
    if multidomain.matrix_dist[i0+1, i1] == 0  # External boundary
        u_block[end, :] .= 0.0  # No-slip wall
        v_block[end, :] .= 0.0
    end
    
    # Bottom boundary (y = y[1])
    if multidomain.matrix_dist[i0, i1-1] == 0  # External boundary
        u_block[:, 1] .= 0.0  # No-slip wall
        v_block[:, 1] .= 0.0
    end
    
    # Top boundary (y = y[end])
    if multidomain.matrix_dist[i0, i1+1] == 0  # External boundary
        # Check if this is a lid boundary (top of domain)
        if i1 == size(multidomain.matrix_dist, 2) - 2  # Top-most active row
            u_block[:, end] .= U_lid  # Moving lid
            v_block[:, end] .= 0.0
        else
            u_block[:, end] .= 0.0  # No-slip wall
            v_block[:, end] .= 0.0
        end
    end
end

"""
    apply_3d_velocity_bc!(u_block::AbstractArray{T,3}, v_block::AbstractArray{T,3}, 
                         w_block::AbstractArray{T,3}, multidomain::MultiDomain{3}, 
                         block_idx::Int, x::Vector, y::Vector, z::Vector, 
                         U_lid::Float64) where T

Apply 3D velocity boundary conditions for a single block.
"""
function apply_3d_velocity_bc!(u_block::AbstractArray{T,3}, v_block::AbstractArray{T,3}, 
                              w_block::AbstractArray{T,3}, multidomain::MultiDomain{3}, 
                              block_idx::Int, x::Vector, y::Vector, z::Vector, 
                              U_lid::Float64) where T
    
    # Get block coordinates in the distribution matrix
    block_coord = multidomain.g[block_idx]
    i0, i1, i2 = block_coord[1], block_coord[2], block_coord[3]
    
    nx, ny, nz = size(u_block)
    
    # Apply boundary conditions for each face
    
    # X- face (x = x[1])
    if multidomain.matrix_dist[i0-1, i1, i2] == 0
        u_block[1, :, :] .= 0.0
        v_block[1, :, :] .= 0.0
        w_block[1, :, :] .= 0.0
    end
    
    # X+ face (x = x[end])  
    if multidomain.matrix_dist[i0+1, i1, i2] == 0
        u_block[end, :, :] .= 0.0
        v_block[end, :, :] .= 0.0
        w_block[end, :, :] .= 0.0
    end
    
    # Y- face (y = y[1])
    if multidomain.matrix_dist[i0, i1-1, i2] == 0
        u_block[:, 1, :] .= 0.0
        v_block[:, 1, :] .= 0.0
        w_block[:, 1, :] .= 0.0
    end
    
    # Y+ face (y = y[end])
    if multidomain.matrix_dist[i0, i1+1, i2] == 0
        # Check if this is a lid boundary
        if i1 == size(multidomain.matrix_dist, 2) - 2  # Top-most active row
            u_block[:, end, :] .= U_lid  # Moving lid in x-direction
            v_block[:, end, :] .= 0.0
            w_block[:, end, :] .= 0.0
        else
            u_block[:, end, :] .= 0.0
            v_block[:, end, :] .= 0.0
            w_block[:, end, :] .= 0.0
        end
    end
    
    # Z- face (z = z[1])
    if multidomain.matrix_dist[i0, i1, i2-1] == 0
        u_block[:, :, 1] .= 0.0
        v_block[:, :, 1] .= 0.0
        w_block[:, :, 1] .= 0.0
    end
    
    # Z+ face (z = z[end])
    if multidomain.matrix_dist[i0, i1, i2+1] == 0
        u_block[:, :, end] .= 0.0
        v_block[:, :, end] .= 0.0
        w_block[:, :, end] .= 0.0
    end
end

"""
    apply_pressure_boundary_conditions!(p::Array, multidomain::MultiDomain{D}, 
                                       x::Vector, y::Vector, z::Union{Vector,Nothing}) where D

Apply pressure boundary conditions for multi-domain setup.
"""
function apply_pressure_boundary_conditions!(p::Array, multidomain::MultiDomain{D}, 
                                            x::Vector, y::Vector, z::Union{Vector,Nothing}) where D
    
    # For incompressible flow, pressure is determined up to a constant
    # Set reference pressure at a specific point to remove nullspace
    
    # Choose the first interior point of the first active block as reference
    if multidomain.index_sum > 0
        # Set reference pressure to zero at interior point
        if D == 2
            p[1, 2, 2] = 0.0  # First block, interior point
        else
            p[1, 2, 2, 2] = 0.0  # First block, interior point
        end
    end
    
    # Apply Neumann boundary conditions (∂p/∂n = prescribed) for external boundaries
    # This is typically handled in the linear solver through appropriate matrix modification
end

"""
    apply_interface_coupling!(u::Array, v::Array, w::Union{Array,Nothing}, 
                             p::Array, multidomain::MultiDomain{D}) where D

Apply coupling conditions at interfaces between domain blocks.
"""
function apply_interface_coupling!(u::Array, v::Array, w::Union{Array,Nothing}, 
                                  p::Array, multidomain::MultiDomain{D}) where D
    
    # Process each pair of adjacent blocks
    for l in 1:multidomain.index_sum
        block_coord = multidomain.g[l]
        
        if D == 2
            i0, i1 = block_coord[1], block_coord[2]
            apply_2d_interface_coupling!(u, v, p, multidomain, l, i0, i1)
        else
            i0, i1, i2 = block_coord[1], block_coord[2], block_coord[3]
            apply_3d_interface_coupling!(u, v, w, p, multidomain, l, i0, i1, i2)
        end
    end
end

"""
    apply_2d_interface_coupling!(u::Array, v::Array, p::Array, 
                                multidomain::MultiDomain{2}, block_idx::Int, 
                                i0::Int, i1::Int)

Apply 2D interface coupling conditions.
"""
function apply_2d_interface_coupling!(u::Array, v::Array, p::Array, 
                                     multidomain::MultiDomain{2}, block_idx::Int, 
                                     i0::Int, i1::Int)
    
    # Check each neighboring block and apply coupling
    neighbors = [(i0-1, i1), (i0+1, i1), (i0, i1-1), (i0, i1+1)]
    directions = [:left, :right, :bottom, :top]
    
    for (idx, (ni0, ni1)) in enumerate(neighbors)
        if (1 <= ni0 <= size(multidomain.matrix_dist, 1) && 
            1 <= ni1 <= size(multidomain.matrix_dist, 2) &&
            multidomain.matrix_dist[ni0, ni1] > 0)
            
            # Find neighbor block index
            neighbor_idx = findfirst(coord -> coord == (ni0, ni1), multidomain.g)
            
            if neighbor_idx !== nothing
                apply_2d_coupling_condition!(u, v, p, block_idx, neighbor_idx, directions[idx])
            end
        end
    end
end

"""
    apply_2d_coupling_condition!(u::Array, v::Array, p::Array, 
                                block1::Int, block2::Int, direction::Symbol)

Apply specific coupling condition between two 2D blocks.
"""
function apply_2d_coupling_condition!(u::Array, v::Array, p::Array, 
                                     block1::Int, block2::Int, direction::Symbol)
    
    # Continuity of velocity and pressure at interfaces
    # This is a simplified version - in practice, would need careful interpolation
    # for different polynomial orders (p-refinement)
    
    if direction == :right
        # Right face of block1 couples with left face of block2
        # Ensure continuity: u1[end, :] = u2[1, :]
        u_avg = 0.5 * (u[block1, end, :] + u[block2, 1, :])
        v_avg = 0.5 * (v[block1, end, :] + v[block2, 1, :])
        p_avg = 0.5 * (p[block1, end, :] + p[block2, 1, :])
        
        u[block1, end, :] .= u_avg
        u[block2, 1, :] .= u_avg
        v[block1, end, :] .= v_avg
        v[block2, 1, :] .= v_avg
        p[block1, end, :] .= p_avg
        p[block2, 1, :] .= p_avg
        
    elseif direction == :top
        # Top face of block1 couples with bottom face of block2
        u_avg = 0.5 * (u[block1, :, end] + u[block2, :, 1])
        v_avg = 0.5 * (v[block1, :, end] + v[block2, :, 1])
        p_avg = 0.5 * (p[block1, :, end] + p[block2, :, 1])
        
        u[block1, :, end] .= u_avg
        u[block2, :, 1] .= u_avg
        v[block1, :, end] .= v_avg
        v[block2, :, 1] .= v_avg
        p[block1, :, end] .= p_avg
        p[block2, :, 1] .= p_avg
    end
    
    # Similar logic for :left and :bottom directions (symmetric)
end

"""
    apply_3d_interface_coupling!(u::Array, v::Array, w::Array, p::Array, 
                                multidomain::MultiDomain{3}, block_idx::Int, 
                                i0::Int, i1::Int, i2::Int)

Apply 3D interface coupling conditions.
"""
function apply_3d_interface_coupling!(u::Array, v::Array, w::Array, p::Array, 
                                     multidomain::MultiDomain{3}, block_idx::Int, 
                                     i0::Int, i1::Int, i2::Int)
    
    # 3D has 6 potential neighbors (±x, ±y, ±z directions)
    neighbors = [(i0-1, i1, i2), (i0+1, i1, i2), 
                (i0, i1-1, i2), (i0, i1+1, i2),
                (i0, i1, i2-1), (i0, i1, i2+1)]
    directions = [:x_minus, :x_plus, :y_minus, :y_plus, :z_minus, :z_plus]
    
    for (idx, (ni0, ni1, ni2)) in enumerate(neighbors)
        if (1 <= ni0 <= size(multidomain.matrix_dist, 1) && 
            1 <= ni1 <= size(multidomain.matrix_dist, 2) &&
            1 <= ni2 <= size(multidomain.matrix_dist, 3) &&
            multidomain.matrix_dist[ni0, ni1, ni2] > 0)
            
            # Find neighbor block index
            neighbor_idx = findfirst(coord -> coord == (ni0, ni1, ni2), multidomain.g)
            
            if neighbor_idx !== nothing
                apply_3d_coupling_condition!(u, v, w, p, block_idx, neighbor_idx, directions[idx])
            end
        end
    end
end

"""
    apply_3d_coupling_condition!(u::Array, v::Array, w::Array, p::Array, 
                                block1::Int, block2::Int, direction::Symbol)

Apply specific coupling condition between two 3D blocks.
"""
function apply_3d_coupling_condition!(u::Array, v::Array, w::Array, p::Array, 
                                     block1::Int, block2::Int, direction::Symbol)
    
    # Similar to 2D but for 3D interfaces
    if direction == :x_plus
        # Right face (x+) of block1 couples with left face (x-) of block2
        u_avg = 0.5 * (u[block1, end, :, :] + u[block2, 1, :, :])
        v_avg = 0.5 * (v[block1, end, :, :] + v[block2, 1, :, :])
        w_avg = 0.5 * (w[block1, end, :, :] + w[block2, 1, :, :])
        p_avg = 0.5 * (p[block1, end, :, :] + p[block2, 1, :, :])
        
        u[block1, end, :, :] .= u_avg
        u[block2, 1, :, :] .= u_avg
        v[block1, end, :, :] .= v_avg
        v[block2, 1, :, :] .= v_avg
        w[block1, end, :, :] .= w_avg
        w[block2, 1, :, :] .= w_avg
        p[block1, end, :, :] .= p_avg
        p[block2, 1, :, :] .= p_avg
        
    # Similar implementations for other directions...
    end
end

"""
    apply_manufactured_boundary_conditions!(u::Array, v::Array, w::Union{Array,Nothing}, 
                                           p::Array, manufactured_sol::NSManufacturedSolution,
                                           multidomain::MultiDomain{D}, x::Vector, y::Vector, 
                                           z::Union{Vector,Nothing}, t::Float64) where D

Apply boundary conditions from manufactured solutions for validation.
"""
function apply_manufactured_boundary_conditions!(u::Array, v::Array, w::Union{Array,Nothing}, 
                                                p::Array, manufactured_sol::NSManufacturedSolution,
                                                multidomain::MultiDomain{D}, x::Vector, y::Vector, 
                                                z::Union{Vector,Nothing}, t::Float64) where D
    
    # Apply exact solution values on boundaries for validation studies
    for l in 1:multidomain.index_sum
        if D == 2
            apply_manufactured_bc_2d!(u[l, :, :], v[l, :, :], p[l, :, :], 
                                     manufactured_sol, x, y, t)
        else
            apply_manufactured_bc_3d!(u[l, :, :, :], v[l, :, :, :], w[l, :, :, :], p[l, :, :, :], 
                                     manufactured_sol, x, y, z, t)
        end
    end
end

"""
    apply_manufactured_bc_2d!(u_block::AbstractMatrix, v_block::AbstractMatrix, 
                              p_block::AbstractMatrix, manufactured_sol::NSManufacturedSolution,
                              x::Vector, y::Vector, t::Float64)

Apply 2D manufactured solution boundary conditions.
"""
function apply_manufactured_bc_2d!(u_block::AbstractMatrix, v_block::AbstractMatrix, 
                                  p_block::AbstractMatrix, manufactured_sol::NSManufacturedSolution,
                                  x::Vector, y::Vector, t::Float64)
    
    nx, ny = size(u_block)
    
    # Apply exact values on all boundaries
    for i in 1:nx
        # Bottom boundary (j=1)
        u_block[i, 1] = calculate_manufactured_u(manufactured_sol, x[i], y[1], 0.0, t)
        v_block[i, 1] = calculate_manufactured_v(manufactured_sol, x[i], y[1], 0.0, t)
        p_block[i, 1] = calculate_manufactured_p(manufactured_sol, x[i], y[1], 0.0, t)
        
        # Top boundary (j=ny)
        u_block[i, ny] = calculate_manufactured_u(manufactured_sol, x[i], y[ny], 0.0, t)
        v_block[i, ny] = calculate_manufactured_v(manufactured_sol, x[i], y[ny], 0.0, t)
        p_block[i, ny] = calculate_manufactured_p(manufactured_sol, x[i], y[ny], 0.0, t)
    end
    
    for j in 1:ny
        # Left boundary (i=1)
        u_block[1, j] = calculate_manufactured_u(manufactured_sol, x[1], y[j], 0.0, t)
        v_block[1, j] = calculate_manufactured_v(manufactured_sol, x[1], y[j], 0.0, t)
        p_block[1, j] = calculate_manufactured_p(manufactured_sol, x[1], y[j], 0.0, t)
        
        # Right boundary (i=nx)
        u_block[nx, j] = calculate_manufactured_u(manufactured_sol, x[nx], y[j], 0.0, t)
        v_block[nx, j] = calculate_manufactured_v(manufactured_sol, x[nx], y[j], 0.0, t)
        p_block[nx, j] = calculate_manufactured_p(manufactured_sol, x[nx], y[j], 0.0, t)
    end
end

"""
    apply_manufactured_bc_3d!(u_block::AbstractArray{T,3}, v_block::AbstractArray{T,3}, 
                              w_block::AbstractArray{T,3}, p_block::AbstractArray{T,3}, 
                              manufactured_sol::NSManufacturedSolution,
                              x::Vector, y::Vector, z::Vector, t::Float64) where T

Apply 3D manufactured solution boundary conditions.
"""
function apply_manufactured_bc_3d!(u_block::AbstractArray{T,3}, v_block::AbstractArray{T,3}, 
                                  w_block::AbstractArray{T,3}, p_block::AbstractArray{T,3}, 
                                  manufactured_sol::NSManufacturedSolution,
                                  x::Vector, y::Vector, z::Vector, t::Float64) where T
    
    nx, ny, nz = size(u_block)
    
    # Apply exact values on all 6 faces
    
    # X faces (i=1 and i=nx)
    for j in 1:ny, k in 1:nz
        # X- face
        u_block[1, j, k] = calculate_manufactured_u(manufactured_sol, x[1], y[j], z[k], t)
        v_block[1, j, k] = calculate_manufactured_v(manufactured_sol, x[1], y[j], z[k], t)
        w_block[1, j, k] = calculate_manufactured_w(manufactured_sol, x[1], y[j], z[k], t)
        p_block[1, j, k] = calculate_manufactured_p(manufactured_sol, x[1], y[j], z[k], t)
        
        # X+ face
        u_block[nx, j, k] = calculate_manufactured_u(manufactured_sol, x[nx], y[j], z[k], t)
        v_block[nx, j, k] = calculate_manufactured_v(manufactured_sol, x[nx], y[j], z[k], t)
        w_block[nx, j, k] = calculate_manufactured_w(manufactured_sol, x[nx], y[j], z[k], t)
        p_block[nx, j, k] = calculate_manufactured_p(manufactured_sol, x[nx], y[j], z[k], t)
    end
    
    # Y faces (j=1 and j=ny)
    for i in 1:nx, k in 1:nz
        # Y- face
        u_block[i, 1, k] = calculate_manufactured_u(manufactured_sol, x[i], y[1], z[k], t)
        v_block[i, 1, k] = calculate_manufactured_v(manufactured_sol, x[i], y[1], z[k], t)
        w_block[i, 1, k] = calculate_manufactured_w(manufactured_sol, x[i], y[1], z[k], t)
        p_block[i, 1, k] = calculate_manufactured_p(manufactured_sol, x[i], y[1], z[k], t)
        
        # Y+ face
        u_block[i, ny, k] = calculate_manufactured_u(manufactured_sol, x[i], y[ny], z[k], t)
        v_block[i, ny, k] = calculate_manufactured_v(manufactured_sol, x[i], y[ny], z[k], t)
        w_block[i, ny, k] = calculate_manufactured_w(manufactured_sol, x[i], y[ny], z[k], t)
        p_block[i, ny, k] = calculate_manufactured_p(manufactured_sol, x[i], y[ny], z[k], t)
    end
    
    # Z faces (k=1 and k=nz)
    for i in 1:nx, j in 1:ny
        # Z- face
        u_block[i, j, 1] = calculate_manufactured_u(manufactured_sol, x[i], y[j], z[1], t)
        v_block[i, j, 1] = calculate_manufactured_v(manufactured_sol, x[i], y[j], z[1], t)
        w_block[i, j, 1] = calculate_manufactured_w(manufactured_sol, x[i], y[j], z[1], t)
        p_block[i, j, 1] = calculate_manufactured_p(manufactured_sol, x[i], y[j], z[1], t)
        
        # Z+ face
        u_block[i, j, nz] = calculate_manufactured_u(manufactured_sol, x[i], y[j], z[nz], t)
        v_block[i, j, nz] = calculate_manufactured_v(manufactured_sol, x[i], y[j], z[nz], t)
        w_block[i, j, nz] = calculate_manufactured_w(manufactured_sol, x[i], y[j], z[nz], t)
        p_block[i, j, nz] = calculate_manufactured_p(manufactured_sol, x[i], y[j], z[nz], t)
    end
end